//
//  Heap.swift
//  Chapter3LinkedList
//
//  Created by TRS-Mobile on 2021/9/2.
//

import Foundation

struct Heap<E: Equatable> {
    var elements: [E] = []
    let sort: (E, E) -> Bool
    init(sort: @escaping (E, E) -> Bool, elements: [E] = []) {
        self.sort = sort
        self.elements = elements
        if !elements.isEmpty {
            for i in stride(from: elements.count / 2 - 1, through: 0, by: -1) {
                siftDown(from: i)
            }
        }
    }
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var count: Int {
        return elements.count
    }
    
    func peek() -> E? {
        return elements.first
    }
    
    func leftChildIndex(ofParentIndex index: Int) -> Int {
        return 2 * index + 1
    }
    
    func rightChildIndex(ofParentIndex index: Int) -> Int {
        return 2 * index + 2
    }
    
    func parentIndex(ofChildAt index: Int) -> Int {
        return (index - 1) / 2
    }
    
    mutating func remove() -> E? {
        guard !isEmpty else {
            return .none
        }
        elements.swapAt(0, count - 1)
        defer {
            siftDown(from: 0)
        }
        return elements.removeLast()
    }
    
    mutating func siftDown(from index: Int) {
//        var parent = index
//        while true {
//            let left = leftChildIndex(ofParentIndex: parent)
//            let right = rightChildIndex(ofParentIndex: parent)
//            var candidate = parent
//            if left < count && sort(elements[left], elements[candidate]) {
//                candidate = left
//            }
//            if right < count && sort(elements[right], elements[candidate]) {
//                candidate = right
//            }
//            if candidate == parent {
//                return
//            }
//            elements.swapAt(candidate, parent)
//            parent = candidate
//        }
        return siftDown(from: index, upTo: count)
    }
    
    mutating func siftDown(from index: Int, upTo size: Int) {
      var parent = index
      while true {
        let left = leftChildIndex(ofParentIndex: parent)
        let right = rightChildIndex(ofParentIndex: parent)
        var candidate = parent
        if left < size && sort(elements[left], elements[candidate]) {
          candidate = left
        }
        if right < size && sort(elements[right], elements[candidate]) {
          candidate = right
        }
        if candidate == parent {
          return
        }
        elements.swapAt(parent, candidate)
        parent = candidate
      }
    }
    
    mutating func insert(_ element: E) {
        elements.append(element)
        siftUp(from: elements.count - 1)
    }
    
    mutating func siftUp(from index: Int) {
        var child = index
        var parent = parentIndex(ofChildAt: child)
        while child > 0 && sort(elements[child], elements[parent]) {
            elements.swapAt(child, parent)
            child = parent
            parent = parentIndex(ofChildAt: child)
        }
    }
    
    mutating func remove(at index: Int) -> E? {
        guard index < elements.count else {
            return .none
        }
        if index == elements.count - 1 {
            return elements.removeLast()
        }
        elements.swapAt(index, elements.count - 1)
        defer {
            siftDown(from: index)
            siftUp(from: index)
        }
        return elements.removeLast()
    }
    
    func index(of element: E, startingAt i: Int) -> Int? {
        if i >= count {
            return .none
        }
        if sort(element, elements[i]) {
            return .none
        }
        if element == elements[i] {
            return i
        }
        if let j = index(of: element, startingAt: leftChildIndex(ofParentIndex: i)) {
            return j
        }
        if let j = index(of: element, startingAt: rightChildIndex(ofParentIndex: i)) {
            return j
        }
        return .none
    }
}
